from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import shutil
from pathlib import Path
import uuid
import pymupdf

app = FastAPI()

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 创建临时文件夹
UPLOAD_DIR = Path("./uploads")
UPLOAD_DIR.mkdir(exist_ok=True)


# def pdf_to_markdown(pdf_path: str) -> str:
#     """将PDF转换成MarkDown"""
#     markdown_text = []
#
#     try:
#         # 打开PDF文件
#         pdf_file = pymupdf.open(pdf_path)
#
#         for page_num in range(pdf_file.page_count):
#             page = pdf_file[page_num]
#
#             # 提取文本
#             text = page.get_text("text")
#
#             # 基本的格式处理
#             blocks = page.get_text("blocks")
#             for block in blocks:
#                 text = block[4]
#                 # 获取字体大小
#                 font_size = block[5]
#
#                 # 根据字体大小确定标题级别
#                 if font_size > 14:
#                     text = f"## {text}\n"
#                 elif font_size > 12:
#                     text = f"### {text}\n"
#                 else:
#                     text = f"{text}\n"
#
#                 markdown_text.append(text)
#
#             markdown_text.append("\n---\n")  # 页面分隔符
#
#         pdf_file.close()
#
#     except Exception as e:
#         raise HTTPException(status_code=500, detail=f"Failed to convert PDF to Markdown: {str(e)}")

def pdf_to_markdown(pdf_path: str) -> str:
    markdown_text = []

    try:
        doc = pymupdf.open(pdf_path)

        for page_num in range(len(doc)):
            page = doc[page_num]

            # 提取图片
            images = page.get_images()
            for img_index, img in enumerate(images):
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]

                # 保存图片
                image_filename = f"image_{page_num}_{img_index}.png"
                image_path = UPLOAD_DIR / image_filename
                with open(image_path, "wb") as image_file:
                    image_file.write(image_bytes)

                # 添加图片引用到Markdown
                markdown_text.append(f"![Image {img_index}]({image_filename})\n")

            # 提取文本
            blocks = page.get_text("blocks")
            for block in blocks:
                text = block[4]
                font_size = block[5]
                font_flags = block[6]  # 获取字体标志

                # 格式处理
                if font_flags & 2 ** 3:  # 粗体
                    text = f"**{text}**"
                if font_flags & 2 ** 1:  # 斜体
                    text = f"*{text}*"

                # 标题处理
                if font_size > 14:
                    text = f"## {text}\n"
                elif font_size > 12:
                    text = f"### {text}\n"
                else:
                    text = f"{text}\n"

                markdown_text.append(text)

            markdown_text.append("\n---\n")

        doc.close()
        return "".join(markdown_text)

    except Exception as e:
        raise Exception(f"PDF转换错误: {str(e)}")


@app.post("/convert")
async def convert_pdf(file: UploadFile = File(...)):
    try:
        # 生成唯一文件名
        filename = f"{uuid.uuid4()}.pdf"
        pdf_path = UPLOAD_DIR / f"{filename}.pdf"

        # 保存上传的pdf
        with pdf_path.open("wb") as f:
            shutil.copyfileobj(file.file, f)

        # 转换pdf到markdown
        markdown_text = pdf_to_markdown(pdf_path)

        # 清理临时文件
        pdf_path.unlink()

        return {"success": True, "content": markdown_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to upload file: {str(e)}")
